// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/specialforms/at_least_n_non_nulls.h>
#include <pollux/common/base/tests/FloatConstants.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/sparksql/tests/SparkFunctionBaseTest.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions::sparksql::test {
namespace {

class AtLeastNNonNullsTest : public SparkFunctionBaseTest {
 public:
  AtLeastNNonNullsTest() {
    // Allow for parsing literal integers as INTEGER, not BIGINT.
    options_.parseIntegerAsBigint = false;
  }
};

TEST_F(AtLeastNNonNullsTest, basic) {
  auto testAtLeastNNonNulls = [&](int32_t n,
                                  const std::vector<VectorPtr>& input,
                                  const VectorPtr& expected) {
    std::string func = fmt::format("at_least_n_non_nulls({}", n);
    for (auto i = 0; i < input.size(); ++i) {
      func += fmt::format(", c{}", i);
    }
    func += ")";
    const auto result = evaluate(func, make_row_vector(input));
    assertEqualVectors(expected, result);
  };
  auto strings = make_nullable_flat_vector<StringView>(
      {std::nullopt, "1", "", std::nullopt, ""});
  auto bools = make_nullable_flat_vector<bool>(
      {std::nullopt, true, false, std::nullopt, std::nullopt});
  auto ints =
      make_nullable_flat_vector<int32_t>({-1, 0, 1, std::nullopt, std::nullopt});
  auto floats = make_nullable_flat_vector<float>(
      {FloatConstants::kMaxF, FloatConstants::kNaNF, 0.1f, 0.0f, std::nullopt});
  auto doubles = make_nullable_flat_vector<double>(
      {std::log(-2.0),
       FloatConstants::kMaxD,
       FloatConstants::kNaND,
       std::nullopt,
       0.1});
  auto arrays = make_array_vector_from_json<int32_t>(
      {"[1, null, 3]", "[1, 2, 3]", "null", "[null]", "[]"});
  auto maps = make_map_vector_from_json<int32_t, int32_t>(
      {"{1: 10, 2: null, 3: null}", "{1: 10, 2: 20}", "{1: 2}", "{}", "null"});
  auto consts = make_constant<int32_t>(2, 5);
  auto indices = make_indices({1, 2, 3, 4, 0});
  auto dicts = wrap_in_dictionary(indices, 5, doubles);

  auto expected = make_flat_vector<bool>({false, true, true, false, false});
  testAtLeastNNonNulls(2, {strings, bools}, expected);

  expected = make_flat_vector<bool>({false, false, false, false, false});
  testAtLeastNNonNulls(3, {strings, bools}, expected);

  expected = make_flat_vector<bool>({true, true, true, true, true});
  testAtLeastNNonNulls(0, {strings, bools}, expected);
  testAtLeastNNonNulls(-1, {strings, bools}, expected);

  expected = make_flat_vector<bool>({true, false, true, true, false});
  testAtLeastNNonNulls(1, {floats}, expected);

  expected = make_flat_vector<bool>({false, true, false, false, true});
  testAtLeastNNonNulls(1, {doubles}, expected);

  expected = make_flat_vector<bool>({false, true, true, false, false});
  testAtLeastNNonNulls(2, {strings, bools, floats}, expected);

  expected = make_flat_vector<bool>({false, true, true, false, false});
  testAtLeastNNonNulls(3, {bools, ints, floats, doubles}, expected);

  expected = make_flat_vector<bool>({false, false, false, false, false});
  testAtLeastNNonNulls(2, {floats, doubles}, expected);

  expected = make_flat_vector<bool>({true, false, false, true, false});
  testAtLeastNNonNulls(4, {maps, arrays, consts, dicts}, expected);
}

TEST_F(AtLeastNNonNullsTest, error) {
  auto input = make_flat_vector<int32_t>({1, 2, 3});

  POLLUX_ASSERT_USER_THROW(
      evaluate("at_least_n_non_nulls(1.0, c0)", make_row_vector({input})),
      "The first input type should be INTEGER but got DOUBLE");
  POLLUX_ASSERT_USER_THROW(
      evaluate("at_least_n_non_nulls(1)", make_row_vector({})),
      "AtLeastNNonNulls expects to receive at least 2 arguments");
  POLLUX_ASSERT_USER_THROW(
      evaluate("at_least_n_non_nulls(c0, c1)", make_row_vector({input, input})),
      "The first parameter should be constant expression");
  POLLUX_ASSERT_USER_THROW(
      evaluate(
          "at_least_n_non_nulls(cast(null as int), c0)",
          make_row_vector({input})),
      "The first parameter should not be null");
}
} // namespace
} // namespace kumo::pollux::functions::sparksql::test
